import random
import heapq
import numpy as np
import time
class n_queen:
    def __init__(self, n):
        self.n = n

    def random_state(self):
        state = [i for i in range(1, self.n+1)]
        random.shuffle(state)
        return state

    def new_state(self, state):
        i, j = random.sample(range(0, self.n), 2)
        new = state.copy()
        new[i], new[j] = state[j], state[i]
        return new

    def check(self, state):
        conflicts = 0
        diag1 = {}
        diag2 = {}
        for i, col in enumerate(state):
            diag_key1 = i - col
            if diag_key1 in diag1:
                conflicts += 1
            diag1[diag_key1] = True

            diag_key2 = i + col
            if diag_key2 in diag2:
                conflicts += 1
            diag2[diag_key2] = True
        return conflicts

    def random_replace(self, tries=100, epoch=10000):
        while tries:
            tries -= 1
            state = self.random_state()
            while epoch:
                epoch -= 1
                states = []
                for i in range(15):
                    _ = self.new_state(state)
                    v = self.check(_)
                    if v < self.check(state):
                        if v == 0:
                            return _
                        heapq.heappush(states, (v, _))
                if len(states)==0:
                    break
                state = heapq.heappop(states)[1]
        return None
    def local_bundle_search(self, k=5, epoch=10000):
        states = [self.random_state() for i in range(k)]
        while epoch:
            epoch -= 1
            new_states = []
            for i in range(k):
                for j in range(10):
                    _s = self.new_state(states[i])
                    if self.check(_s) == 0:
                        return _s
                    heapq.heappush(new_states, (self.check(_s), _s))
            states = heapq.nsmallest(k, new_states)
            states = [i[1] for i in states]
        return None
    def random_bundle_search(self, k=5, epoch=10000):
        states = [self.random_state() for i in range(k)]
        while epoch:
            epoch -= 1
            new_states = []
            v = []
            for i in range(k):
                for j in range(10):
                    _s = self.new_state(states[i])
                    if self.check(_s) == 0:
                        return _s
                    new_states.append(_s)
                    v.append(self.check(_s))
            v = np.array(v)
            v = 1 - (v - v.min())/(v.max() - v.min())
            states = random.choices(new_states, weights=v, k=k)
        return None
    def show(self, state):
        state = np.array(state) - 1
        state = np.eye(8, dtype=int)[state]
        for i in range(state.shape[0]):
            for j in range(state.shape[1]):
                print(state[i][j], end=' ')
            print()

    def test(self, epoch = 1000):
        funcs = [self.random_replace, self.local_bundle_search, self.random_bundle_search]
        for func in funcs:
            st = time.time()
            for i in range(epoch):
                state = func()
            et = time.time()
            print(func.__name__, "消耗时间为：", et - st)
            if state==None:
                print("未得到解")
            else:
                n_queen.show(state)


if __name__ == '__main__':
    n_queen = n_queen(8)
    n_queen.test()
